"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/10/11 16:51
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    gen_psu_color_disp_image.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    gen_psu_color_disp_image_2024/10/11.0.1

-------------------------------------------------------------------------------
# @Description:

    ---------------------------------------------------------------------------
    [Update History]:
        2024/10/11:

-------------------------------------------------------------------------------
"""
import os
import cv2
import numpy as np
from utils import Utils, Colorizer


class GenFalseColorImage(object):
    def __init__(self, levels=4096, depth_range=(0, 3000)):
        self.colorize = Colorizer(levels=levels, depth_range=depth_range)

    def gen_false_color_image(self, src_img_path: str, dst_img_path: str, min_dis=10, max_dis=3000) -> None:
        """

        :param src_img_path: depth image path
        :param dst_img_path:
        :param min_dis: valid depth range min distance
        :param max_dis: valid depth range max distance
        :return:
        """
        if not os.path.exists(src_img_path):
            raise FileNotFoundError(src_img_path)

        img = cv2.imread(src_img_path, cv2.IMREAD_UNCHANGED)
        img = np.clip(img, min_dis, max_dis)
        pse_img = self.colorize.colorize(img, (0, 3000), bins=6000)
        cv2.imwrite(dst_img_path, pse_img)


if __name__ == '__main__':
    root_path = '/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800'
    dp_img_path = 'complex-geometry/curved-surfaces/ball/0003_indoor_mix_pu_depth.png'

    full_dp_img_path = os.path.join(root_path, dp_img_path)
    full_ps_img_path = full_dp_img_path.replace('.png', '_pse.jpg')

    gen_pse = GenFalseColorImage()
    gen_pse.gen_false_color_image(full_dp_img_path, full_ps_img_path)
